# -*- coding: utf-8 -*-
"""
数据库连接重试工具
@api-version: 1.1.0
// REF: API-DB-RETRY-001
"""

import time
import logging
from functools import wraps
from sqlalchemy.exc import DisconnectionError, OperationalError, TimeoutError
from backend.models import db

logger = logging.getLogger(__name__)

class DatabaseRetryError(Exception):
    """数据库重试失败异常"""
    pass

def db_retry(max_retries=3, delay=0.1, backoff=2.0):
    """
    数据库操作重试装饰器
    
    Args:
        max_retries (int): 最大重试次数
        delay (float): 初始延迟时间（秒）
        backoff (float): 延迟时间倍数
    
    Returns:
        装饰器函数
    """
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            last_exception = None
            
            for attempt in range(max_retries + 1):
                try:
                    return func(*args, **kwargs)
                    
                except (DisconnectionError, OperationalError, TimeoutError) as e:
                    last_exception = e
                    
                    # 记录重试日志
                    logger.warning(
                        f"数据库操作失败，第{attempt + 1}次尝试: {str(e)}"
                    )
                    
                    # 如果是最后一次尝试，直接抛出异常
                    if attempt == max_retries:
                        logger.error(
                            f"数据库操作重试{max_retries}次后仍然失败: {str(e)}"
                        )
                        break
                    
                    # 回滚当前事务
                    try:
                        db.session.rollback()
                    except Exception as rollback_error:
                        logger.warning(f"事务回滚失败: {str(rollback_error)}")
                    
                    # 等待后重试
                    sleep_time = delay * (backoff ** attempt)
                    logger.info(f"等待{sleep_time:.2f}秒后重试...")
                    time.sleep(sleep_time)
                    
                except Exception as e:
                    # 非连接相关异常，直接抛出
                    logger.error(f"数据库操作出现非连接异常: {str(e)}")
                    raise e
            
            # 所有重试都失败，抛出自定义异常
            raise DatabaseRetryError(
                f"数据库操作重试{max_retries}次后失败: {str(last_exception)}"
            )
        
        return wrapper
    return decorator

def execute_with_retry(operation, max_retries=3, delay=0.1, backoff=2.0):
    """
    执行数据库操作并自动重试
    
    Args:
        operation (callable): 要执行的数据库操作函数
        max_retries (int): 最大重试次数
        delay (float): 初始延迟时间（秒）
        backoff (float): 延迟时间倍数
    
    Returns:
        操作结果
    
    Raises:
        DatabaseRetryError: 重试失败异常
    """
    last_exception = None
    
    for attempt in range(max_retries + 1):
        try:
            return operation()
            
        except (DisconnectionError, OperationalError, TimeoutError) as e:
            last_exception = e
            
            logger.warning(
                f"数据库操作失败，第{attempt + 1}次尝试: {str(e)}"
            )
            
            if attempt == max_retries:
                logger.error(
                    f"数据库操作重试{max_retries}次后仍然失败: {str(e)}"
                )
                break
            
            # 回滚当前事务
            try:
                db.session.rollback()
            except Exception as rollback_error:
                logger.warning(f"事务回滚失败: {str(rollback_error)}")
            
            # 等待后重试
            sleep_time = delay * (backoff ** attempt)
            logger.info(f"等待{sleep_time:.2f}秒后重试...")
            time.sleep(sleep_time)
            
        except Exception as e:
            # 非连接相关异常，直接抛出
            logger.error(f"数据库操作出现非连接异常: {str(e)}")
            raise e
    
    # 所有重试都失败
    raise DatabaseRetryError(
        f"数据库操作重试{max_retries}次后失败: {str(last_exception)}"
    )

def check_db_connection():
    """
    检查数据库连接状态
    
    Returns:
        bool: 连接是否正常
    """
    try:
        # 执行简单查询测试连接
        db.session.execute('SELECT 1')
        return True
    except Exception as e:
        logger.error(f"数据库连接检查失败: {str(e)}")
        return False

def get_connection_pool_status():
    """
    获取连接池状态信息
    
    Returns:
        dict: 连接池状态信息
    """
    try:
        engine = db.engine
        pool = engine.pool
        
        return {
            'pool_size': pool.size(),
            'checked_in': pool.checkedin(),
            'checked_out': pool.checkedout(),
            'overflow': pool.overflow(),
            'invalid': pool.invalid()
        }
    except Exception as e:
        logger.error(f"获取连接池状态失败: {str(e)}")
        return {}

# 使用示例：
# @db_retry(max_retries=3)
# def create_user(user_data):
#     user = User(**user_data)
#     db.session.add(user)
#     db.session.commit()
#     return user

# 或者：
# result = execute_with_retry(lambda: User.query.all())